package spark_sql;

import bean.AvgBuffer;
import bean.User;
import java.util.Arrays;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.TypedColumn;
import org.apache.spark.sql.expressions.Aggregator;


/**
 * @author shihb
 * @date 2020/1/10 22:56
 * 用户自定义强类型聚合函数
 */
public class UdafClassDemo {

  public static void main(String[] args) {
    //local模式,创建SparkConf对象设定spark的部署环境
    SparkConf sparkConf = new SparkConf().setMaster("local[*]").setAppName("");
    //创建spark上下文对象
    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
    //创建SparkSQL的环境对象
    SparkSession spark = new SparkSession(jsc.sc());




    //创建Rdd
    JavaRDD<User> rdd = jsc.parallelize(Arrays.asList(
        new User(1, "zhangsan", 20),
        new User(2, "lisi", 30),
        new User(3, "wangwu", 40)
    ),1);


    Encoder<User> userEncoder = Encoders.bean(User.class);
    Dataset<User> dataset = spark.createDataset(rdd.rdd(), userEncoder);
    dataset.createOrReplaceTempView("user");
    

    //创建udaf函数
    MyAvgAgeClassFunction udaf = new MyAvgAgeClassFunction();
    //讲强类型的udaf转为列类型
    TypedColumn<User, Double> avgAgeCol = udaf.toColumn();
    //使用DSL来调用强类型的udaf
    dataset.select(avgAgeCol).show();





    //释放资源
    spark.stop();
  }

}

/**
 * 用户自定义聚合函数强类型实现
 * 求age平均值
 */
class MyAvgAgeClassFunction extends Aggregator<User, AvgBuffer,Double>{

  @Override
  public AvgBuffer zero() {
    return new AvgBuffer(0L,0L);
  }

  @Override
  public AvgBuffer reduce(AvgBuffer avgBuffer, User user) {
    return new AvgBuffer(avgBuffer.getSum()+user.getAge(),avgBuffer.getCount()+1);
  }

  @Override
  public AvgBuffer merge(AvgBuffer ab1, AvgBuffer ab2) {
    return new AvgBuffer(ab1.getSum()+ab2.getSum(),ab1.getCount()+ab2.getCount());
  }

  @Override
  public Double finish(AvgBuffer avgBuffer) {
    return new Double((avgBuffer.getSum()+0.0)/avgBuffer.getCount());
  }

  @Override
  public Encoder<AvgBuffer> bufferEncoder() {
    return Encoders.bean(AvgBuffer.class);
  }

  @Override
  public Encoder<Double> outputEncoder() {
    return Encoders.DOUBLE();
  }
}